from cryptography import x509
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.asymmetric import rsa
from cryptography.hazmat.primitives.serialization import Encoding, PrivateFormat, NoEncryption, BestAvailableEncryption
from cryptography.x509.oid import NameOID
from cryptography.hazmat.primitives import serialization
from cryptography.hazmat.backends import default_backend
import datetime
import ipaddress  # 导入 ipaddress 模块

# 生成 CA 密钥对
ca_key = rsa.generate_private_key(
    public_exponent=65537,
    key_size=2048,
    backend=default_backend()
)

# 生成 CA 证书
ca_subject = x509.Name([
    x509.NameAttribute(NameOID.COUNTRY_NAME, u"US"),
    x509.NameAttribute(NameOID.STATE_OR_PROVINCE_NAME, u"California"),
    x509.NameAttribute(NameOID.LOCALITY_NAME, u"San Francisco"),
    x509.NameAttribute(NameOID.ORGANIZATION_NAME, u"My CA"),
    x509.NameAttribute(NameOID.COMMON_NAME, u"My CA"),
])

ca_cert = x509.CertificateBuilder(
    subject_name=ca_subject,
    issuer_name=ca_subject,
    public_key=ca_key.public_key(),
    serial_number=x509.random_serial_number(),
    not_valid_before=datetime.datetime.utcnow(),
    not_valid_after=datetime.datetime.utcnow() + datetime.timedelta(days=365),
).add_extension(
    x509.SubjectKeyIdentifier.from_public_key(ca_key.public_key()),
    critical=False,
).add_extension(
    x509.AuthorityKeyIdentifier.from_issuer_public_key(ca_key.public_key()),
    critical=False,
).add_extension(
    x509.BasicConstraints(ca=True, path_length=None),
    critical=True,
).sign(
    private_key=ca_key,
    algorithm=hashes.SHA256(),
    backend=default_backend()
)

# 保存 CA 证书和密钥
with open("ca.crt", "wb") as f:
    f.write(ca_cert.public_bytes(Encoding.PEM))

with open("ca.key", "wb") as f:
    f.write(ca_key.private_bytes(
        encoding=Encoding.PEM,
        format=PrivateFormat.TraditionalOpenSSL,
        encryption_algorithm=NoEncryption()
    ))

# 生成服务器密钥对
server_key = rsa.generate_private_key(
    public_exponent=65537,
    key_size=2048,
    backend=default_backend()
)

# 生成服务器证书
server_subject = x509.Name([
    x509.NameAttribute(NameOID.COUNTRY_NAME, u"US"),
    x509.NameAttribute(NameOID.STATE_OR_PROVINCE_NAME, u"California"),
    x509.NameAttribute(NameOID.LOCALITY_NAME, u"San Francisco"),
    x509.NameAttribute(NameOID.ORGANIZATION_NAME, u"My Organization"),
    x509.NameAttribute(NameOID.COMMON_NAME, u"myserver.com"),
])

server_cert = x509.CertificateBuilder(
    subject_name=server_subject,
    issuer_name=ca_subject,
    public_key=server_key.public_key(),
    serial_number=x509.random_serial_number(),
    not_valid_before=datetime.datetime.utcnow(),
    not_valid_after=datetime.datetime.utcnow() + datetime.timedelta(days=365),
).add_extension(
    x509.SubjectKeyIdentifier.from_public_key(server_key.public_key()),
    critical=False,
).add_extension(
    x509.AuthorityKeyIdentifier.from_issuer_public_key(ca_key.public_key()),
    critical=False,
).add_extension(
    x509.BasicConstraints(ca=False, path_length=None),
    critical=True,
).add_extension(
    x509.SubjectAlternativeName([
        x509.DNSName(u"myserver.com"),
        x509.IPAddress(ipaddress.ip_address(u"10.0.200.89"))
    ]),
    critical=False,
).sign(
    private_key=ca_key,
    algorithm=hashes.SHA256(),
    backend=default_backend()
)

# 保存服务器证书和密钥
with open("server.crt", "wb") as f:
    f.write(server_cert.public_bytes(Encoding.PEM))

with open("server.key", "wb") as f:
    f.write(server_key.private_bytes(
        encoding=Encoding.PEM,
        format=PrivateFormat.TraditionalOpenSSL,
        encryption_algorithm=NoEncryption()
    ))
